from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
from matplotlib import pyplot as plt
from matplotlib.pyplot import *
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas


class Calculation_Stokes(QWidget):
	IsStokes_signal = pyqtSignal()

	def __init__(self):
		super(Calculation_Stokes, self).__init__()
		self.setWindowTitle("")
		# self.setWindowModality(Qt.WindowModal)
		self.resize(550, 454)  # 设置控件大小
		self.setStyleSheet("""QWidget{color:black;}""")

		self.all_layout = QVBoxLayout()
		self.up_biggest_box = QGroupBox()
		self.up_biggest_box.setTitle("计算Stokes矢量")
		self.all_layout.addWidget(self.up_biggest_box)

		self.up_biggest_layout = QHBoxLayout()
		self.up_biggest_box.setLayout(self.up_biggest_layout)

		self.up_biggest_box = QGroupBox()
		self.up_biggest_layout.addWidget(self.up_biggest_box)

		self.up_all_layout = QHBoxLayout()
		self.up_biggest_box.setLayout(self.up_all_layout)

		self.first_layout = QHBoxLayout()
		self.up_all_layout.addLayout(self.first_layout)

		self.first_box = QGroupBox("振幅相位法")
		self.first_box.setCheckable(True)
		self.first_box.setChecked(False)
		self.label11 = QLabel("振幅Ex:")
		self.label12 = QLabel("振幅Ey:")
		self.label13 = QLabel("相位差:")
		self.label14 = QLabel("角度")
		self.linedit11 = QLineEdit()    # 振幅ExEdit
		self.linedit12 = QLineEdit()    # 振幅EyEdit
		self.linedit13 = QLineEdit()    # 相位差Edit

		self.first_box_layout = QGridLayout()
		self.first_box_layout.addWidget(self.label11, 0, 0, 1, 1)
		self.first_box_layout.addWidget(self.label12, 1, 0, 1, 1)
		self.first_box_layout.addWidget(self.label13, 2, 0, 1, 1)
		self.first_box_layout.addWidget(self.linedit11, 0, 1, 1, 1)
		self.first_box_layout.addWidget(self.linedit12, 1, 1, 1, 1)
		self.first_box_layout.addWidget(self.linedit13, 2, 1, 1, 1)
		self.first_box_layout.addWidget(self.label14, 2, 2, 1, 1)

		self.first_box.setLayout(self.first_box_layout)
		self.first_layout.addWidget(self.first_box)

		self.two_layout = QHBoxLayout()
		self.up_all_layout.addLayout(self.two_layout)

		self.two_box = QGroupBox("偏振角度法")
		self.two_box.setCheckable(True)
		self.two_box.setChecked(True)

		self.checkbox1 = QCheckBox("自然光")
		self.checkbox1.setChecked(True)
		self.checkbox2 = QCheckBox("左旋圆偏光")
		self.checkbox3 = QCheckBox("右旋圆偏光")
		self.checkbox4 = QCheckBox("线偏光:")

		self.linedit21 = QLineEdit()    # 偏振角度法-输入角度
		self.linedit21.setReadOnly(True)
		self.label21 = QLabel("角度")
		self.two_box_layout = QGridLayout()
		self.two_box_layout.addWidget(self.checkbox1, 0, 0, 1, 4)

		self.two_box_layout.addWidget(self.checkbox2, 1, 0, 1, 4)
		self.two_box_layout.addWidget(self.checkbox3, 2, 0, 1, 4)
		self.two_box_layout.addWidget(self.checkbox4, 3, 0, 1, 1)
		self.two_box_layout.addWidget(self.linedit21, 3, 1, 1, 1)
		self.two_box_layout.addWidget(self.label21, 3, 2, 1, 1)

		self.two_box.setLayout(self.two_box_layout)
		self.two_layout.addWidget(self.two_box)

		self.up_button_layout = QVBoxLayout()
		self.action_calculate_button = QPushButton("开\n始\n计\n算")
		self.action_calculate_button.setMaximumWidth(30)
		self.copy_result = QPushButton("复\n制\n结\n果")
		self.copy_result.setMaximumWidth(30)

		self.up_button_layout.addWidget(self.action_calculate_button)
		self.up_button_layout.addWidget(self.copy_result)
		self.up_biggest_layout.addLayout(self.up_button_layout)

		self.result_v_layout = QVBoxLayout()
		self.result_layout = QGridLayout()
		self.label31 = QLabel("归一化Stokes矢量")
		self.label32 = QLabel("I")
		self.label33 = QLabel("Q")
		self.label34 = QLabel("U")
		self.label35 = QLabel("V")
		self.linedit31 = QLineEdit()
		self.linedit32 = QLineEdit()
		self.linedit33 = QLineEdit()
		self.linedit34 = QLineEdit()
		self.result_layout.addWidget(self.label31, 0, 0, 1, 3)
		self.result_layout.addWidget(self.label32, 1, 0, 1, 1)
		self.result_layout.addWidget(self.label33, 2, 0, 1, 1)
		self.result_layout.addWidget(self.label34, 3, 0, 1, 1)
		self.result_layout.addWidget(self.label35, 4, 0, 1, 1)
		self.result_layout.addWidget(self.linedit31, 1, 1, 1, 1)
		self.result_layout.addWidget(self.linedit32, 2, 1, 1, 1)
		self.result_layout.addWidget(self.linedit33, 3, 1, 1, 1)
		self.result_layout.addWidget(self.linedit34, 4, 1, 1, 1)
		self.result_v_layout.addLayout(self.result_layout)
		self.up_biggest_layout.addLayout(self.result_v_layout)

		self.middle_all_layout = QHBoxLayout()
		self.all_layout.addLayout(self.middle_all_layout)

		self.middle_all_left_layout = QVBoxLayout()
		self.middle_all_layout.addLayout(self.middle_all_left_layout)

		self.middle_biggest_box = QGroupBox()
		self.middle_biggest_box.setTitle("反演偏振状态")
		self.middle_biggest_layout = QGridLayout()
		self.middle_biggest_box.setLayout(self.middle_biggest_layout)
		self.middle_all_left_layout.addWidget(self.middle_biggest_box)

		self.over_result_layout = QGridLayout()
		self.middle_biggest_layout.addLayout(self.over_result_layout, 0, 0, 1, 1)
		self.label41 = QLabel("I")
		self.label42 = QLabel("Q")
		self.label43 = QLabel("U")
		self.label44 = QLabel("V")
		self.linedit41 = QLineEdit()
		self.linedit42 = QLineEdit()
		self.linedit43 = QLineEdit()
		self.linedit44 = QLineEdit()
		self.over_result_layout.addWidget(self.label41, 0, 0, 1, 1)
		self.over_result_layout.addWidget(self.label42, 1, 0, 1, 1)
		self.over_result_layout.addWidget(self.label43, 2, 0, 1, 1)
		self.over_result_layout.addWidget(self.label44, 3, 0, 1, 1)
		self.over_result_layout.addWidget(self.linedit41, 0, 1, 1, 1)
		self.over_result_layout.addWidget(self.linedit42, 1, 1, 1, 1)
		self.over_result_layout.addWidget(self.linedit43, 2, 1, 1, 1)
		self.over_result_layout.addWidget(self.linedit44, 3, 1, 1, 1)

		self.middle_box1_layout = QHBoxLayout()
		self.middle_biggest_layout.addLayout(self.middle_box1_layout, 0, 1, 1, 1)
		self.middle_box1 = QGroupBox("振幅相位法")
		self.label51 = QLabel("总偏振度")
		self.label52 = QLabel("线偏振度")
		self.label53 = QLabel("圆偏振度")
		self.linedit51 = QLineEdit()
		self.linedit52 = QLineEdit()
		self.linedit53 = QLineEdit()

		self.middle_box1_g_layout = QGridLayout()
		self.middle_box1_g_layout.addWidget(self.label51, 0, 0, 1, 1)
		self.middle_box1_g_layout.addWidget(self.label52, 1, 0, 1, 1)
		self.middle_box1_g_layout.addWidget(self.label53, 2, 0, 1, 1)
		self.middle_box1_g_layout.addWidget(self.linedit51, 0, 1, 1, 1)
		self.middle_box1_g_layout.addWidget(self.linedit52, 1, 1, 1, 1)
		self.middle_box1_g_layout.addWidget(self.linedit53, 2, 1, 1, 1)

		self.middle_box1.setLayout(self.middle_box1_g_layout)
		self.middle_box1_layout.addWidget(self.middle_box1)

		self.middle_box2_layout = QHBoxLayout()
		self.middle_biggest_layout.addLayout(self.middle_box2_layout, 0, 2, 1, 1)

		self.middle_box2 = QGroupBox("完全偏振光参数")
		self.label61 = QLabel("振幅Ex:")
		self.label62 = QLabel("振幅Ey:")
		self.label63 = QLabel("相位差:")

		self.linedit61 = QLineEdit()
		self.linedit62 = QLineEdit()
		self.linedit63 = QLineEdit()
		self.linedit63.setPlaceholderText("角度")

		self.middle_box2_g_layout = QGridLayout()
		self.middle_box2_g_layout.addWidget(self.label61, 0, 0, 1, 1)
		self.middle_box2_g_layout.addWidget(self.label62, 1, 0, 1, 1)
		self.middle_box2_g_layout.addWidget(self.label63, 2, 0, 1, 1)
		self.middle_box2_g_layout.addWidget(self.linedit61, 0, 1, 1, 1)
		self.middle_box2_g_layout.addWidget(self.linedit62, 1, 1, 1, 1)
		self.middle_box2_g_layout.addWidget(self.linedit63, 2, 1, 1, 1)

		self.middle_box2.setLayout(self.middle_box2_g_layout)
		self.middle_box2_layout.addWidget(self.middle_box2)

		self.action_over_button = QPushButton("开始反演")
		self.middle_biggest_layout.addWidget(self.action_over_button, 1, 0, 1, 1)

		self.middle_box3_layout = QHBoxLayout()
		self.middle_biggest_layout.addLayout(self.middle_box3_layout, 1, 1, 1, 2)

		self.middle_box3 = QGroupBox("偏振光类型描述")

		self.linedit71 = QLineEdit()

		self.middle_box3_g_layout = QHBoxLayout()
		self.middle_box3_g_layout.addWidget(self.linedit71)

		self.middle_box3.setLayout(self.middle_box3_g_layout)
		self.middle_box3_layout.addWidget(self.middle_box3)

		self.clear_button = QPushButton("清空")
		self.clear_button.setFixedWidth(50)
		self.middle_all_left_layout.addWidget(self.clear_button)

		self.middle_right_layout = QVBoxLayout()
		self.middle_all_layout.addLayout(self.middle_right_layout)

		self.fig1 = figure(figsize=(3, 3))
		self.canvas1 = FigureCanvas(self.fig1)
		self.middle_right_layout.addWidget(self.canvas1)
		self.ax = self.fig1.subplots()
		self.ax.spines['top'].set_visible(False)
		self.ax.spines['right'].set_visible(False)
		self.ax.spines['bottom'].set_position(('data', 0))
		self.ax.spines['left'].set_position(('axes', 0.5))
		self.ax.set_aspect('equal')

		self.middle_right_below_layout = QHBoxLayout()
		self.middle_right_layout.addLayout(self.middle_right_below_layout)

		self.draw_button = QPushButton("绘图")
		self.save_button = QPushButton("保存")
		self.middle_right_below_layout.addWidget(self.draw_button)
		self.middle_right_below_layout.addWidget(self.save_button)

		self.setLayout(self.all_layout)

		self.I_angle = None
		self.Q_angle = None
		self.U_angle = None
		self.V_angle = None
		self.setup_signal()

	def setup_signal(self):
		self.first_box.clicked.connect(self.checked_box1)   # 相位法
		self.two_box.clicked.connect(self.checked_box2)     # 偏振角度法
		self.checkbox1.stateChanged.connect(self.check1)    # 偏振角度算法---子选项的互斥
		self.checkbox2.stateChanged.connect(self.check2)
		self.checkbox3.stateChanged.connect(self.check3)
		self.checkbox4.stateChanged.connect(self.check4)
		self.action_calculate_button.clicked.connect(self.action_calculate_function)    # 计算
		self.copy_result.clicked.connect(self.copy_result2over)             # 复制结果
		self.action_over_button.clicked.connect(self.over_result)           # 反演
		self.draw_button.clicked.connect(self.draw_plot)                    # 画图
	# 相位法
	def checked_box1(self):
		if self.first_box.isChecked():
			self.linedit21.clear()
			self.two_box.setChecked(False)
	# 偏振角度法
	def checked_box2(self):
		if self.two_box.isChecked():
			self.first_box.setChecked(False)
			self.linedit11.clear()
			self.linedit12.clear()
			self.linedit13.clear()
	# 偏振角度算法---子选项的互斥
	def check1(self, status):
		if status:
			self.checkbox2.setChecked(False)
			self.checkbox3.setChecked(False)
			self.checkbox4.setChecked(False)
	def check2(self, status):
		if status:
			self.checkbox1.setChecked(False)
			self.checkbox3.setChecked(False)
			self.checkbox4.setChecked(False)
	def check3(self, status):
		if status:
			self.checkbox1.setChecked(False)
			self.checkbox2.setChecked(False)
			self.checkbox4.setChecked(False)
	def check4(self, status):
		if status:
			self.checkbox1.setChecked(False)
			self.checkbox2.setChecked(False)
			self.checkbox3.setChecked(False)
			self.linedit21.setReadOnly(False)
		else:
			self.linedit21.clear()
			self.linedit21.setReadOnly(True)
	# 计算
	def action_calculate_function(self):
		if self.first_box.isChecked():
			if self.linedit11.text() == "" or self.linedit12.text() == "" or self.linedit13.text() == "":
				print(QMessageBox.warning(self, '：', '请检查输入参数是否准确！', QMessageBox.Yes, QMessageBox.Yes))
			else:
				angle = float(self.linedit13.text())
				self.I_angle = 1.000000
				self.Q_angle = 0.000000
				self.U_angle = np.cos(angle * np.pi / 180)
				self.V_angle = np.sin(angle * np.pi / 180)
				I0 = '{:.6f}'.format(self.I_angle)
				Q0 = '{:.6f}'.format(self.Q_angle)
				U0 = '{:.6f}'.format(self.U_angle)
				V0 = '{:.6f}'.format(self.V_angle)
				self.linedit31.setText(I0)
				self.linedit32.setText(Q0)
				self.linedit33.setText(U0)
				self.linedit34.setText(V0)

		elif self.two_box.isChecked():
			I0 = '{:.6f}'.format(1.0)
			Q0 = '{:.6f}'.format(0.0)
			U0 = '{:.6f}'.format(0.0)
			V0 = '{:.6f}'.format(0.0)
			V1 = '{:.6f}'.format(-1.0)
			V2 = '{:.6f}'.format(1.0)
			if not self.checkbox4.isChecked():
				if self.checkbox1.isChecked():
					self.linedit31.setText(I0)
					self.linedit32.setText(Q0)
					self.linedit33.setText(U0)
					self.linedit34.setText(V0)
				elif self.checkbox2.isChecked():
					self.linedit31.setText(I0)
					self.linedit32.setText(Q0)
					self.linedit33.setText(U0)
					self.linedit34.setText(V1)
				elif self.checkbox3.isChecked():
					self.linedit31.setText(I0)
					self.linedit32.setText(Q0)
					self.linedit33.setText(U0)
					self.linedit34.setText(V2)
				else:
					pass

				self.I_angle = float(self.linedit31.text())
				self.Q_angle = float(self.linedit32.text())
				self.U_angle = float(self.linedit33.text())
				self.V_angle = float(self.linedit34.text())
			else:
				if self.linedit21.text() is not None:
					phi1 = float(self.linedit21.text())
					z = 1
					a1 = np.abs(z * np.cos(phi1 * np.pi / 180))
					a2 = np.abs(z * np.sin(phi1 * np.pi / 180))

					phi2 = 360 - phi1
					deta = phi2 - phi1
					S0 = np.power(a1, 2) + np.power(a2, 2)

					S1 = np.power(a1, 2) - np.power(a2, 2)

					S2 = a1 * a2 * np.exp(-0.0j * deta) + a1 * a2 * np.exp(0.0j * deta)
					S2 = np.real(S2)
					if 90 < phi1 < 180 or 270 < phi1 < 360:
						S2 = -S2
					else:
						S2 = S2

					S3 = (a1 * a2 * np.exp(0.0j * deta) - a1 * a2 * np.exp(-0.0j * deta)) * 0.0j
					S3 = np.real(S3)

					self.I_angle = '{:.6f}'.format(S0)
					self.Q_angle = '{:.6f}'.format(S1)
					self.U_angle = '{:.6f}'.format(S2)
					self.V_angle = '{:.6f}'.format(S3)
					if S0 >= 0.99999:
						S0 = 1.0
					if S1 >= 0.99999:
						S1 = 1.0
					if np.abs(S0) <= 0.000001:
						S0 = 0
					if np.abs(S1) <= 0.000001:
						S1 = 0
					if np.abs(S2) <= 0.000001:
						S2 = 0
					if np.abs(S3) <= 0.000001:
						S3 = 0

					S0 = '{:.6f}'.format(S0)
					S1 = '{:.6f}'.format(S1)
					S2 = '{:.6f}'.format(S2)
					S3 = '{:.6f}'.format(S3)
					self.linedit31.setText(S0)
					self.linedit32.setText(S1)
					self.linedit33.setText(S2)
					self.linedit34.setText(S3)
				else:
					print(QMessageBox.warning(self, '：', '请检查输入参数是否准确！', QMessageBox.Yes, QMessageBox.Yes))
		else:
			pass
	# 复制结果
	def copy_result2over(self):
		I = self.linedit31.text()
		Q = self.linedit32.text()
		U = self.linedit33.text()
		V = self.linedit34.text()
		if I == "":
			print(QMessageBox.warning(self, '：', '请检查输入参数是否准确！', QMessageBox.Yes, QMessageBox.Yes))
		else:
			self.linedit41.setText(I)
			self.linedit42.setText(Q)
			self.linedit43.setText(U)
			self.linedit44.setText(V)
	# 反演
	def over_result(self):
		I = float('{:.4f}'.format(float(self.I_angle)))
		Q = float('{:.4f}'.format(float(self.Q_angle)))
		U = float('{:.4f}'.format(float(self.U_angle)))
		V = float('{:.4f}'.format(float(self.V_angle)))
		if self.two_box.isChecked():
			if self.checkbox1.isChecked():
				self.linedit51.setText(str(0))
				self.linedit52.setText(str(0))
				self.linedit53.setText(str(0))
			elif self.checkbox2.isChecked() or self.checkbox3.isChecked():
				self.linedit51.setText(str(1.000))
				self.linedit52.setText(str(0.000))
				self.linedit53.setText(str(1.000))
			else:
				all_dop = 1.0000
				y_dop = all_dop - I
				'{:.3f}'.format(all_dop)
				self.linedit51.setText(str('{:.3f}'.format(all_dop)))
				self.linedit52.setText(str('{:.3f}'.format(I)))
				self.linedit53.setText(str('{:.3f}'.format(y_dop)))
		elif self.first_box.isChecked():
			all_dop = 1.0000
			angle1 = float(self.linedit13.text())
			x_dop = all_dop * np.cos(angle1 * np.pi / 180)
			y_dop = all_dop * np.sin(angle1 * np.pi / 180)
			'{:.3f}'.format(all_dop)
			self.linedit51.setText(str('{:.3f}'.format(all_dop)))
			self.linedit52.setText(str('{:.3f}'.format(x_dop)))
			self.linedit53.setText(str('{:.3f}'.format(y_dop)))
		else:
			all_dop = 1.0000
			y_dop = all_dop - I
			self.linedit51.setText(str(all_dop))
			self.linedit52.setText(str(I))
			self.linedit53.setText(str(y_dop))
		if self.linedit13.text() != "" or self.linedit21.text() != "":
			if self.linedit13.text() != "":
				phi1 = float(self.linedit13.text())
				self.linedit61.setText(self.linedit11.text())
				self.linedit62.setText(self.linedit12.text())
				self.linedit63.setText(str(phi1))
			else:
				phi1 = float(self.linedit21.text())
				if phi1 <= 90:
					self.linedit63.setText(str(0.0))
				else:
					self.linedit63.setText(str(180.0))
				x_f = np.cos(phi1 * np.pi / 180)
				y_f = np.sin(phi1 * np.pi / 180)
				x_f = abs(x_f)
				x_f = '{:.3f}'.format(x_f)
				y_f = abs(y_f)
				y_f = '{:.3f}'.format(y_f)
				self.linedit61.setText(str(x_f))
				self.linedit62.setText(str(y_f))
		else:
			if self.checkbox1.isChecked():
				self.linedit61.setText("未知")
				self.linedit62.setText("未知")
				self.linedit63.setText("未知")
			elif self.checkbox2.isChecked():
				self.linedit61.setText("0.707")
				self.linedit62.setText("0.707")
				self.linedit63.setText("-90")
			elif self.checkbox3.isChecked():
				self.linedit61.setText("0.707")
				self.linedit62.setText("0.707")
				self.linedit63.setText("90")
		if self.two_box.isChecked():
			if self.checkbox1.isChecked():
				self.linedit71.setText("自然光")
			elif self.checkbox2.isChecked():
				self.linedit71.setText("左旋圆偏振光")
			elif self.checkbox3.isChecked():
				self.linedit71.setText("右旋圆偏振光")
			elif self.checkbox4.isChecked():
				if self.linedit21.text() == "0" or self.linedit21.text() == "-0" or self.linedit21.text() == "180" or self.linedit21.text() == "-180":
					self.linedit71.setText("水平线偏振光")
				elif self.linedit21.text() == "90" or self.linedit21.text() == "-90":
					self.linedit71.setText("垂直线偏振光")
				else:
					angle = float(self.linedit21.text())

					if angle >= 0:
						self.linedit71.setText(str(angle) + "°线偏振光")
					else:
						x = 180 - angle
						self.linedit71.setText(str(x) + "°线偏振光")
			else:
				pass
		elif self.first_box.isChecked():
			angle = float(self.linedit13.text())
			if angle > 360:
				x = int(angle / 360)
				angle = angle - x * 360
			else:
				angle = angle
			if 90 > angle >= 0:
				self.linedit71.setText("右旋椭圆偏振光(长轴偏转45.0°)")
			elif angle == 90:
				self.linedit71.setText("右旋圆偏振光")
			elif 180 > angle > 90:
				self.linedit71.setText("右旋椭圆偏振光(长轴偏转-45.0度)")
			elif angle == 180:
				self.linedit71.setText("135度线偏振光")
			elif 270 > angle > 180:
				self.linedit71.setText("左旋椭圆偏振光(长轴偏转-45.0度)")
			elif angle == 270:
				self.linedit71.setText("左旋圆偏振光")
			elif 360 > angle > 270:
				self.linedit71.setText("左旋椭圆偏振光(长轴偏转45.0度)")
			elif angle == 360:
				self.linedit71.setText("45度线偏振光")
			else:
				pass
	# 画图
	def draw_plot(self):
		plt.close(self.fig1)
		plt.cla()
		circle = plt.Circle((0, 0), radius=0.68, facecolor='white', edgecolor='black', alpha=1.0)
		if self.two_box.isChecked():
			if self.checkbox1.isChecked():
				pass
			elif self.checkbox2.isChecked() or self.checkbox3.isChecked():
				x1 = np.arange(-1, 1, 0.0001)
				y1 = np.sqrt(1 - np.power((x1 - 2), 2)) + 2
				x2 = np.arange(-1, 1, 0.0001)
				y2 = -1 * np.sqrt(1 - np.power((x2 - 2), 2)) + 2
				self.ax.add_artist(circle)
				self.ax.plot(x1, y1, x2, y2, color='k', linestyle=':')
				self.canvas1.draw()
			else:
				x = float(self.linedit61.text())
				y = float(self.linedit62.text())
				theta = np.arctan(x / y) * 180 / np.pi
				x_zuo = np.arange(-1, 1, 0.1)
				y_zuo = x_zuo * (1 / np.tan(theta * np.pi / 180))
				self.ax.add_artist(circle)
				self.ax.plot(x_zuo, y_zuo, color='k', linestyle=':')
				self.canvas1.draw()


if __name__ == '__main__':
	from sys import argv, exit

	app = QApplication(argv)
	calculate_stokes = Calculation_Stokes()
	calculate_stokes.show()
	exit(app.exec_())
